package com.chenfan.common.extension.lock;

import com.chenfan.common.extension.spring.TransactionUtils;
import org.apache.commons.lang3.StringUtils;
import org.springframework.data.redis.core.StringRedisTemplate;
import org.springframework.util.Assert;

import java.util.Objects;
import java.util.Optional;
import java.util.UUID;
import java.util.concurrent.TimeUnit;
import java.util.function.Supplier;

/**
 * 分布式锁 Redis分布式锁助手（默认锁有效期1分钟）
 *
 * @author wulg
 * @date 2021-08-16
 **/
public class DistributedLockTemplate {

    /**
     * 默认睡眠时间（单位：毫秒）
     */
    private final long DEFAULT_SLEEP_TIME_MS = 100L;

    private final StringRedisTemplate redisTemplate;

    /**
     * Redis分布式锁助手
     */
    private static Optional<DistributedLockTemplate> lockOptional = Optional.empty();

    private static final Supplier<DistributedLockException> NOT_CONFIGURED_BEAN_SUPPLIER = () -> new DistributedLockException(DistributedLockException.Type.NOT_CONFIG_BEAN);

    private static final Supplier<DistributedLockException> NOT_CONFIG_DEPEND_SUPPLIER = () -> new DistributedLockException(DistributedLockException.Type.NOT_CONFIG_DEPEND);

    private static final DistributedLockException GET_LOCK_FAIL_EXP = new DistributedLockException(DistributedLockException.Type.GET_LOCK_FAIL);

    private static Supplier<RuntimeException> GET_LOCK_FAIL_SUPPLIER = () -> GET_LOCK_FAIL_EXP;

    private static final String MACHINE_ID = UUID.randomUUID().toString();

    /**
     * 默认锁有效时间10秒（单位：毫秒）
     */
    private final long lockExpireMs;


    public DistributedLockTemplate(StringRedisTemplate redisTemplate) {
        this(redisTemplate, null);
    }

    public DistributedLockTemplate(StringRedisTemplate redisTemplate, RuntimeException getLockException) {
        this(redisTemplate, getLockException, 10L, TimeUnit.SECONDS);
    }

    public DistributedLockTemplate(StringRedisTemplate redisTemplate, RuntimeException getLockException, long lockExpire, TimeUnit timeUnit) {
        this.redisTemplate = Optional.ofNullable(redisTemplate).orElseThrow(NOT_CONFIG_DEPEND_SUPPLIER);
        lockOptional = Optional.of(this);
        if(Objects.nonNull(getLockException)) {
            GET_LOCK_FAIL_SUPPLIER = () -> getLockException;
        }
        this.lockExpireMs =  timeUnit.toMillis(lockExpire);
    }



    /**
     * 分布式锁处理模板执行器
     *
     * @param lockKey        分布式锁key
     * @param resultSupplier 分布式锁处理回调
     * @return
     */
    public static <T> T execute(String lockKey, Supplier<T> resultSupplier) {
        return execute(lockKey, resultSupplier, GET_LOCK_FAIL_SUPPLIER,5 * 1000L, TimeUnit.MILLISECONDS);
    }

    /**
     * 分布式锁处理模板执行器
     *
     * @param lockKey        分布式锁key
     * @param resultSupplier 分布式锁处理回调
     * @return
     */
    public static <T> T execute(String lockKey, Supplier<T> resultSupplier, Supplier<RuntimeException> exceptionSupplier) {
        return execute(lockKey, resultSupplier, exceptionSupplier,5 * 1000L, TimeUnit.MILLISECONDS);
    }

    /**
     * 分布式锁处理模板执行器
     *
     * @param lockKey        分布式锁key
     * @param resultSupplier 分布式锁处理回调
     * @param waitTimeout    锁等待超时时间
     * @param waitUnit       时间单位
     * @return
     */
    public static <T> T execute(String lockKey, Supplier<T> resultSupplier, long waitTimeout, TimeUnit waitUnit) {
        return execute(lockKey, resultSupplier, GET_LOCK_FAIL_SUPPLIER, waitTimeout, waitUnit);
    }


    /**
     * 分布式锁处理模板执行器
     *
     * @param lockKey           分布式锁key
     * @param resultSupplier    业务务处理结果回调
     * @param waitTimeout       锁等待超时时间
     * @param waitUnit          时间单位
     * @param getLockFailSupplier 获取分布式锁失败异常回调
     * @return
     */
    public static <T> T execute(String lockKey, Supplier<T> resultSupplier, Supplier<RuntimeException> getLockFailSupplier, long waitTimeout, TimeUnit waitUnit) {
        Assert.isTrue(StringUtils.isNotBlank(lockKey), "锁Key不能为空");
        Assert.notNull(resultSupplier, "业务结果处理回调不能为空");
        Assert.isTrue(waitTimeout > 0L, "锁等待超时时间不能为空");
        Assert.notNull(waitUnit, "锁等待超时时间单位不能为空");
        Assert.notNull(getLockFailSupplier, "获取分布式锁失败异常回调不能为空");
        int locked = 0;
        try {
            locked = lockOptional.orElseThrow(NOT_CONFIGURED_BEAN_SUPPLIER).tryLock(lockKey, waitUnit.toMillis(waitTimeout));
            if (locked == 0) {
                throw getLockFailSupplier.get();
            }
            return resultSupplier.get();
        } finally {
            if(locked == 1) {
                lockOptional.orElseThrow(NOT_CONFIGURED_BEAN_SUPPLIER).unlock(lockKey);
            }
        }
    }

    /**
     * @param lockKey       锁键
     * @param waitTimeoutMs 锁有效期（单位：毫秒）
     * @return 0-失败、1-成功、2-重入
     */
    private int tryLock(String lockKey, long waitTimeoutMs) {
        return this.tryLock(lockKey, lockExpireMs, waitTimeoutMs);
    }

    /**
     * 尝试锁，可参考redisson 锁重入做计数
     *
     * @param lockKey       锁的键
     * @param expireTimeMs  锁有效期（单位：毫秒）
     * @param waitTimeoutMs 锁等待超时时间（单位：毫秒）
     * @return 0-失败、1-成功、2-重入
     */
    private int tryLock(String lockKey, long expireTimeMs, long waitTimeoutMs) {
        while (waitTimeoutMs > 0L) {
            String lockValue = this.getLockValue();
            if (redisTemplate.opsForValue().setIfAbsent(lockKey, lockValue, expireTimeMs, TimeUnit.MILLISECONDS).booleanValue()) {
                return 1;
            }
            //允许锁重入，避免嵌套重复加锁导致死锁
            String targetValue = redisTemplate.opsForValue().get(lockKey);
            if (lockValue.equals(targetValue)) {
                return 2;
            }
            long sleepTime;
            if (waitTimeoutMs > DEFAULT_SLEEP_TIME_MS) {
                sleepTime = DEFAULT_SLEEP_TIME_MS;
                waitTimeoutMs -= DEFAULT_SLEEP_TIME_MS;
            } else {
                sleepTime = waitTimeoutMs;
                waitTimeoutMs = 0L;
            }
            try {
                TimeUnit.MILLISECONDS.sleep(sleepTime);
            } catch (InterruptedException e) {
                if (Thread.currentThread().isInterrupted()) {
                    //线程中断恢复
                    Thread.currentThread().interrupt();
                }
                //线程中断后跳出取锁循环
                break;
            }
        }
        return 0;
    }

    /**
     * 解锁
     */
    private void unlock(String lockKey) {
        String targetValue = redisTemplate.opsForValue().get(lockKey);
        // 1、比较锁线程ID，避免B线程获取到的锁被之前A线程释放
        // 2、分布式锁续期问题 -可以参照redisson看门狗机制
        if (!this.getLockValue().equals(targetValue)) {
            return;
        }
        //规避顺序问题导致事务未结束就释放锁
        TransactionUtils.afterCompletion(() -> redisTemplate.delete(lockKey));
    }

    /**
     * 获取锁值
     * @return
     */
    private String getLockValue() {
        return StringUtils.joinWith(":", MACHINE_ID, String.valueOf(Thread.currentThread().getId()));
    }
}